import getpass
import os
from langchain.chat_models import init_chat_model
from langchain_core.messages import HumanMessage, AIMessage
from langgraph.checkpoint.memory import MemorySaver
from langgraph.graph import START, MessagesState, StateGraph
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from typing import Sequence
from langchain_core.messages import BaseMessage
from langgraph.graph.message import add_messages
from typing_extensions import Annotated, TypedDict

os.environ["LANGSMITH_TRACING"] = "false"
#os.environ["LANGSMITH_API_KEY"] = getpass.getpass()
#os.environ["ANTHROPIC_API_KEY"] = "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"
os.environ["OPENAI_API_KEY"] = "sk-dgmoK2W2ajUXlwEE2RxbViNVdri2GAyAh6Ma4wYxfeYppUSc"

if not os.environ.get("OPENAI_API_KEY"):
    os.environ["OPENAI_API_KEY"] = getpass.getpass("Enter API key for OpenAI: ")

model = init_chat_model("deepseek-r1", model_provider="openai", base_url="https://api.lkeap.cloud.tencent.com/v1")

prompt_template = ChatPromptTemplate.from_messages(
    [
        (
            "system",
            "You are a helpful assistant. Answer all questions to the best of your ability in {language}.",
        ),
        MessagesPlaceholder(variable_name="messages"),
    ]
)

class State(TypedDict):
    messages: Annotated[Sequence[BaseMessage], add_messages]
    language: str

def call_model(state: State):
    prompt = prompt_template.invoke(state)
    response = model.invoke(prompt)
    return {"messages": [response]}

workflow = StateGraph(state_schema=State)
workflow.add_edge(START, "model")
workflow.add_node("model", call_model)

memory = MemorySaver()
app = workflow.compile(checkpointer=memory)

config = {"configurable": {"thread_id": "abc789"}}
query = "Hi I'm Todd, please tell me a joke."
language = "English"

input_messages = [HumanMessage(query)]
for chunk, metadata in app.stream(
        {"messages": input_messages, "language": language},
        config,
        stream_mode="messages",
):
    if isinstance(chunk, AIMessage):  # Filter to just model responses
        print(chunk.content, end="|")
